import torch
import torch.nn as nn

from .models.backbone import backbone_factory
from .models.head import YoloHead


class YoloModel(nn.Module):
    def __init__(
        self, backbone: str, grid_size: int, bounding_box_num: int, class_num: int
    ) -> None:
        super().__init__()
        self.grid_size = grid_size
        self.bounding_box_num = bounding_box_num
        self.class_num = class_num
        self.backbone = backbone_factory(backbone)
        self.head = YoloHead(
            backbone=self.backbone,
            grid_size=grid_size,
            bounding_box_num=bounding_box_num,
            class_num=class_num,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.head(self.backbone(x))

    @torch.no_grad
    def postprocess(self, output: torch.Tensor) -> torch.Tensor:
        for i in range(self.bounding_box_num):
            output[..., 0 + i * 5] /= self.grid_size
            output[..., 1 + i * 5] /= self.grid_size
        for i in range(self.grid_size):
            for j in range(self.grid_size):
                for k in range(self.bounding_box_num):
                    output[i, j, 0 + k * 5] += i * (1 / self.grid_size)
                    output[i, j, 1 + k * 5] += j * (1 / self.grid_size)
        return output

    @torch.no_grad
    def predict(
        self,
        x: torch.Tensor,
        confidence_threshold: float = 0.1,
        iou_threshold: float = 0.3,
    ) -> torch.Tensor:
        # outpush shape = [1, gs, gs, 5 * bbn + class_num]
        output = self.forward(x).squeeze(0)
        output = self.postprocess(output)
        output = self.nms(output, confidence_threshold, iou_threshold)
        return output.numpy()

    @torch.no_grad
    def nms(
        self, output: torch.Tensor, confidence_threshold: float, iou_threshold: float
    ):
        # output shape = [gs * gs,5 * bbn + class_num]
        bbox_prob = (
            output[..., 5 * self.bounding_box_num :]
            .reshape((-1, self.class_num))
            .repeat_interleave(2, dim=0)
        )
        bbox_info = output[..., : 5 * self.bounding_box_num].reshape((-1, 5))

        confidence = bbox_info[:, 4].unsqueeze(1)
        bbox_prob = bbox_prob * confidence
        bbox_prob[bbox_prob < confidence_threshold] = 0
        bbox_info = bbox_info[:, :4]

        for i in range(self.class_num):
            rank = torch.argsort(bbox_prob[:, i], descending=True)
            for j in range(bbox_info.shape[0]):
                if bbox_prob[rank[j], i] == 0:
                    continue
                for k in range(j + 1, bbox_info.shape[0]):
                    if bbox_prob[rank[k], i] == 0:
                        continue
                    if self.iou(bbox_info[rank[j]], bbox_info[rank[k]]) > iou_threshold:
                        bbox_prob[rank[k], i] = 0

        filter = torch.max(bbox_prob, dim=1).values > 0
        bbox_info = bbox_info[filter]
        bbox_prob = bbox_prob[filter]
        if bbox_info.shape[0] == 0:
            return torch.tensor([])
        rtn = torch.zeros((bbox_info.shape[0], 6))
        rtn[:, :4] = bbox_info
        rtn[:, 4] = torch.max(bbox_prob, dim=1).values
        rtn[:, 5] = torch.argmax(bbox_prob, dim=1)
        return rtn

    @torch.no_grad
    def iou(self, area1: torch.Tensor, area2: torch.Tensor) -> float:
        # area [center_x, center_y, width, height]
        bbox1 = torch.zeros_like(area1)
        bbox1[0] = area1[0] - area1[2] / 2
        bbox1[1] = area1[1] - area1[3] / 2
        bbox1[2] = area1[0] + area1[2] / 2
        bbox1[3] = area1[1] + area1[3] / 2

        bbox2 = torch.zeros_like(area2)
        bbox2[0] = area2[0] - area2[2] / 2
        bbox2[1] = area2[1] - area2[3] / 2
        bbox2[2] = area2[0] + area2[2] / 2
        bbox2[3] = area2[1] + area2[3] / 2

        # bbox [x1, y1, x2, y2]
        wi = torch.min(bbox1[2], bbox2[2]) - torch.max(bbox1[0], bbox2[0])
        hi = torch.min(bbox1[3], bbox2[3]) - torch.max(bbox1[1], bbox2[1])

        wi = torch.clamp(wi, min=0)
        hi = torch.clamp(hi, min=0)

        intersection = wi * hi
        union = (
            (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
            + (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
            - intersection
        )

        return (intersection / (union + 1e-6)).item()
